"""
滑窗分类器实现类

这个模块提供了滑窗分类检测的具体实现，包括：
1. 224x224像素滑窗，步长224
2. 模糊检测(类别0)
3. 边界框绘制和结果可视化
4. 批量处理功能

注意：这是内部实现类，用户应该通过src.interface.SlidingWindowClassifier来使用。
"""

import os
import sys
from pathlib import Path
from typing import List, Dict, Any, Optional, Tuple
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import cv2

# 添加项目根目录到Python路径
project_root = str(Path(__file__).parent.parent.parent)
if project_root not in sys.path:
    sys.path.insert(0, project_root)

from devdeploy.inference.inference import create_inference


class SlidingWindowClassifierImpl:
    """
    滑窗分类器类
    
    提供滑窗检测功能，主要用于检测图像中的模糊区域。
    使用224x224像素的滑窗，步长224，对检测到的模糊区域(类别0)绘制边界框。
    """
    
    def __init__(self, 
                 model_path: str, 
                 device: str = 'cpu',
                 window_size: int = 224,
                 stride: int = 224,
                 blur_class_id: int = 0,
                 confidence_threshold: float = 0.5):
        """
        初始化滑窗分类器
        
        Args:
            model_path (str): 分类模型文件路径
            device (str): 推理设备，'cpu' 或 'cuda'
            window_size (int): 滑窗大小，默认224
            stride (int): 滑窗步长，默认224
            blur_class_id (int): 模糊类别ID，默认0
            confidence_threshold (float): 置信度阈值，默认0.5
        """
        self.model_path = model_path
        self.device = device
        self.window_size = window_size
        self.stride = stride
        self.blur_class_id = blur_class_id
        self.confidence_threshold = confidence_threshold
        
        # 创建推理器
        self.infer = create_inference(model_path, device)
        
        # 支持的图像格式
        self.supported_formats = ('.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif')
        
        print(f"滑窗分类器初始化完成: {window_size}x{window_size}, 步长{stride}, 设备{device}")
    
    def _get_sliding_windows(self, image: Image.Image) -> List[Tuple[int, int, Image.Image]]:
        """
        生成滑窗
        
        Args:
            image (Image.Image): 输入图像
            
        Returns:
            List[Tuple[int, int, Image.Image]]: 滑窗列表，每个元素为(x, y, window_image)
        """
        width, height = image.size
        windows = []
        
        # 计算滑窗位置
        for y in range(0, height - self.window_size + 1, self.stride):
            for x in range(0, width - self.window_size + 1, self.stride):
                # 裁剪窗口
                window = image.crop((x, y, x + self.window_size, y + self.window_size))
                windows.append((x, y, window))
        
        # 处理边界情况：如果图像尺寸不能被步长整除，添加边界窗口
        # 右边界
        if width % self.stride != 0:
            for y in range(0, height - self.window_size + 1, self.stride):
                x = width - self.window_size
                window = image.crop((x, y, x + self.window_size, y + self.window_size))
                windows.append((x, y, window))
        
        # 下边界
        if height % self.stride != 0:
            for x in range(0, width - self.window_size + 1, self.stride):
                y = height - self.window_size
                window = image.crop((x, y, x + self.window_size, y + self.window_size))
                windows.append((x, y, window))
        
        # 右下角
        if width % self.stride != 0 and height % self.stride != 0:
            x = width - self.window_size
            y = height - self.window_size
            window = image.crop((x, y, x + self.window_size, y + self.window_size))
            windows.append((x, y, window))
        
        return windows
    
    def _detect_blur_in_image(self, image: Image.Image) -> List[Dict[str, Any]]:
        """
        检测单张图像中的模糊区域
        
        Args:
            image (Image.Image): 输入图像
            
        Returns:
            List[Dict[str, Any]]: 模糊区域列表，每个元素包含坐标和置信度
        """
        blur_regions = []
        
        # 生成滑窗
        windows = self._get_sliding_windows(image)
        
        # print(f"  生成 {len(windows)} 个滑窗")
        
        # 对每个窗口进行分类
        for i, (x, y, window) in enumerate(windows):
            try:
                # 使用推理器进行分类
                result = self.infer.infer_single(window, return_prob=True)
                
                # 检查是否为模糊类别且置信度超过阈值
                if (result['class_id'] == self.blur_class_id and 
                    result['confidence'] >= self.confidence_threshold):
                    
                    blur_regions.append({
                        'bbox': [x, y, x + self.window_size, y + self.window_size],
                        'confidence': result['confidence'],
                        'class_id': result['class_id'],
                        'class_name': result['class_name']
                    })
                    
            except Exception as e:
                # print(f"  警告：窗口 {i+1} 推理失败: {str(e)}")
                continue
        
        return blur_regions
    
    def _draw_boxes_on_image(self, image: Image.Image, blur_regions: List[Dict[str, Any]]) -> Image.Image:
        """
        在图像上绘制边界框
        
        Args:
            image (Image.Image): 原始图像
            blur_regions (List[Dict[str, Any]]): 模糊区域列表
            
        Returns:
            Image.Image: 绘制了边界框的图像
        """
        # 创建图像副本
        result_image = image.copy()
        draw = ImageDraw.Draw(result_image)
        
        # 设置绘制参数
        box_color = (255, 0, 0)  # 红色
        text_color = (255, 255, 255)  # 白色
        box_width = 2
        
        # 尝试加载字体
        try:
            font = ImageFont.truetype("/System/Library/Fonts/Arial.ttf", 16)
        except:
            try:
                font = ImageFont.load_default()
            except:
                font = None
        
        # 绘制每个模糊区域
        for i, region in enumerate(blur_regions):
            x1, y1, x2, y2 = region['bbox']
            confidence = region['confidence']
            
            # 绘制边界框
            draw.rectangle([x1, y1, x2, y2], outline=box_color, width=box_width)
            
            # 绘制标签
            label = f"模糊 {confidence:.2f}"
            if font:
                # 获取文本尺寸
                bbox = draw.textbbox((0, 0), label, font=font)
                text_width = bbox[2] - bbox[0]
                text_height = bbox[3] - bbox[1]
                
                # 绘制文本背景
                text_bg = [x1, y1 - text_height - 4, x1 + text_width + 4, y1]
                draw.rectangle(text_bg, fill=box_color)
                
                # 绘制文本
                draw.text((x1 + 2, y1 - text_height - 2), label, fill=text_color, font=font)
            else:
                # 使用默认字体
                draw.text((x1, y1 - 20), label, fill=text_color)
        
        return result_image
    
    def detect_blur_single(self, image, output_path: Optional[str] = None) -> Dict[str, Any]:
        """
        检测单张图片中的模糊区域
        
        Args:
            image: 图像数据，可以是图像路径(str)、PIL Image对象或numpy数组
            output_path (str, optional): 输出图片路径，如果提供则保存带标注的图片
            
        Returns:
            Dict[str, Any]: 检测结果
        """
        try:
            import numpy as np
            
            # 处理输入图像
            if isinstance(image, str):
                # 如果是路径，使用skimage.io.imread加载图像
                from skimage import io
                image_array = io.imread(image)
                # 转换为PIL Image对象
                if image_array.ndim == 3:
                    image_obj = Image.fromarray(image_array)
                else:
                    image_obj = Image.fromarray(image_array).convert('RGB')
                image_path = image
            elif isinstance(image, Image.Image):
                # 如果是PIL Image对象，直接使用
                image_obj = image.convert('RGB')
                image_path = "PIL_Image"
            elif isinstance(image, np.ndarray):
                # 如果是numpy数组，处理单通道和3通道格式
                if len(image.shape) == 2:
                    # 单通道图像（灰度图）
                    if image.dtype != np.uint8:
                        image = (image * 255).astype(np.uint8)
                    # 将单通道转换为3通道RGB
                    image_rgb = np.stack([image, image, image], axis=2)
                    image_obj = Image.fromarray(image_rgb)
                    image_path = "numpy_array_single_channel"
                elif len(image.shape) == 3 and image.shape[2] == 3:
                    # 3通道图像（RGB或BGR）
                    if image.dtype != np.uint8:
                        image = (image * 255).astype(np.uint8)
                    # 假设是BGR格式，转换为RGB
                    image_rgb = image[:, :, ::-1]  # BGR -> RGB
                    image_obj = Image.fromarray(image_rgb)
                    image_path = "numpy_array_3_channel"
                else:
                    raise ValueError(f"不支持的numpy数组形状: {image.shape}，期望形状为(H, W)或(H, W, 3)")
            else:
                raise ValueError(f"不支持的图像输入类型: {type(image)}")
            
            # 检测模糊区域
            blur_regions = self._detect_blur_in_image(image_obj)
            
            # 准备结果
            result = {
                'image_path': image_path,
                'image_size': image_obj.size,
                'blur_regions': blur_regions,
                'total_blur_regions': len(blur_regions)
            }
            
            # 如果提供了输出路径，保存带标注的图像
            if output_path:
                if blur_regions:
                    annotated_image = self._draw_boxes_on_image(image_obj, blur_regions)
                else:
                    annotated_image = image_obj
                
                # 确保输出目录存在
                os.makedirs(os.path.dirname(output_path), exist_ok=True)
                annotated_image.save(output_path)
                result['output_path'] = output_path
            
            return result
            
        except Exception as e:
            error_msg = f"检测图片失败: {str(e)}"
            return {
                'image_path': str(image) if not isinstance(image, str) else image,
                'error': error_msg,
                'blur_regions': [],
                'total_blur_regions': 0
            }
    

    
    def get_model_info(self) -> Dict[str, Any]:
        """
        获取模型信息
        
        Returns:
            Dict[str, Any]: 模型信息
        """
        return self.infer.get_model_info()
    
    def validate_model(self) -> bool:
        """
        验证模型是否有效
        
        Returns:
            bool: 模型是否有效
        """
        return self.infer.validate_model()
